{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ssseg.modules import *\n",
    "import os\n",
    "import torch\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = '0,1,2,3'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_cgf = dict(\n",
    "    train = dict(\n",
    "        type='vspw',\n",
    "        set='train',\n",
    "        rootdir='/data/yixiang/Dataset/VSPW_480p',\n",
    "        aug_opts=[\n",
    "            ('Resize', {'output_size': (720, 480), 'keep_ratio': True, 'scale_range': (0.5, 2.0)}),\n",
    "                     ('RandomCrop', {'crop_size': (512, 512), 'one_category_max_ratio': 0.75}),\n",
    "                     ('RandomFlip', {'flip_prob': 0.5}),\n",
    "                     ('PhotoMetricDistortion', {}),\n",
    "                     ('Normalize', {'mean': [123.675, 116.28, 103.53], 'std': [58.395, 57.12, 57.375]}),\n",
    "                     ('ToTensor', {}),\n",
    "                     ('Padding', {'output_size': (512, 512), 'data_type': 'tensor'})\n",
    "        ],\n",
    "        clip_num=4,\n",
    "        dilation=\"3,6,9\",\n",
    "        random_select=False,\n",
    "        sequence_range=0,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = BuildDataset(mode='TRAIN', logger_handle=Logger('test.log'), dataset_cfg=data_cgf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "124"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset.num_classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader_cfg = dict(\n",
    "    train = dict(\n",
    "        type = ['nondistributed', 'distributed'][0],\n",
    "        batch_size = 8,\n",
    "        num_workers = 1,\n",
    "        shuffle = True,\n",
    "        pin_memory = True,\n",
    "        drop_last = True,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataloader = BuildParallelDataloader(mode='TRAIN', dataset=dataset, cfg=dataloader_cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "x=iter(dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_CFG = dict(\n",
    "    type = 'clip_psp',\n",
    "    num_classes = 124,\n",
    "    benchmark = True,\n",
    "    is_multi_gpus = False,\n",
    "    align_corners = False,\n",
    "    psp_weight = False,\n",
    "    deep_sup_scale = True,\n",
    "    distributed = dict(\n",
    "        is_on = False, \n",
    "        backend = 'nccl'),\n",
    "    norm_cfg = dict(\n",
    "        type ='batchnorm2d', \n",
    "        opts = {}),\n",
    "    act_cfg = dict(\n",
    "        type = 'relu',\n",
    "        opts = dict(inplace=True)),\n",
    "    backbone = dict( \n",
    "        type = 'swin_base_patch4_window12_384_22k',\n",
    "        series = 'swin',\n",
    "        pretrained = True,\n",
    "        selected_indices = (0, 1, 2, 3),\n",
    "        pretrained_model_path = '/home/lja/pretrain/swin_base_patch4_window12_384_22k.pth',\n",
    "        norm_cfg = {'type': 'layernorm', 'opts': {}},\n",
    "    ),\n",
    "    ppm = dict(\n",
    "        in_channels = 2048,\n",
    "        out_channels = 124,\n",
    "        pool_scales = [1, 2, 3, 6],\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BuildModel(cfg=MODEL_CFG, mode='TRAIN')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torch.nn.DataParallel(model)\n",
    "model = model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "b = x.next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 3, 512, 512])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b['image'].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n"
     ]
    }
   ],
   "source": [
    "loss,acc = model(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss.mean().backward()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "40d3a090f54c6569ab1632332b64b2c03c39dcf918b08424e98f38b5ae0af88f"
  },
  "kernelspec": {
   "display_name": "Python 3.8.5 64-bit ('base': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
